#pragma once

#include <string>
#include <vector>
#include <iostream>
#include <fstream>
#include <boost/filesystem.hpp>
#include <openvino/openvino.hpp>
// 加密库 cryptopp
#include "aes.h"
#include "modes.h"
#include "filters.h"

#include <filesystem>

/**
 * @brief 加密头部结构，用于存储文件大小信息
 */
struct st_encrypt_header {
    size_t struct_size;    // 加密信息结构体大小
    size_t xml_size;       // XML文件大小
    size_t bin_size;       // BIN文件大小
};

// 锚点和步长的最大数量定义
static constexpr int MAX_NUM_ANCHORS = 4;
static constexpr int MAX_STRIDES = 4;

/**
 * @brief OpenVINO加密信息结构体，用于存储模型相关信息
 */
struct st_encrypt_info_ovino {
    char gpu_model[100] = {0};    // GPU型号

    int batch_size = 0;           // 批处理大小
    int class_num = 0;            // 类别数量
    int input_channel = 0;        // 输入通道数
    int image_width = 0;          // 图像宽度
    int image_height = 0;         // 图像高度

    int num_anchors_size = 0;     // 锚点数量
    int num_anchors[MAX_NUM_ANCHORS] = {0};  // 各类型锚点数量

    int stride_num = 0;           // 步长数量
    int strides[MAX_STRIDES] = {0};  // 不同步长值

    bool dynamic = false;         // 是否动态分配
    int engine_size = 0;          // 引擎大小

    float obj_threshold{};        // 目标检测阈值
    float nms_threshold{};        // 非极大值抑制阈值
};

/**
 * @brief 精度模式枚举
 */
enum PRECISION_MODE
{
    PRECISION_FP32,    // 32位浮点数精度
    PRECISION_FP16,    // 16位浮点数精度
    PRECISION_INT8,    // 8位整数精度
};

/**
 * @brief 模型转换输入结构体
 */
struct st_trans_model_input
{
    std::string cfg_path;     // 配置文件路径
    std::string model_path;   // 模型文件路径

    PRECISION_MODE precision = PRECISION_FP16;  // 默认使用FP16精度
    std::string gpu_model = "Intel core";      // 使用的GPU型号

    std::string trans_path;   // 转换后的文件路径
};

namespace fs = boost::filesystem;

/**
 * @brief YOLO v6 OpenVINO实现类
 */
class yolo_v6_o {
public:
    yolo_v6_o();
    ~yolo_v6_o();

    /**
     * @brief 将YOLO v6模型文件转换为OpenVINO IR并加密
     * @param input 包含配置文件路径等信息的转换输入结构体
     * @return 转换成功返回true，否则返回false
     */
    bool TransModel(const st_trans_model_input& input);
};
